import argparse
import json
import os
import sys
import ipe_param_ocr
from ctypes import *
import struct

sys.path.append("../../../")

from common.python.common import *
from common.python.callback_data_struct import *
from common.python.infer_process import *
from common.python import image_decoder
from common.python import image_encoder

pythonapi.PyCapsule_GetPointer.restype = c_void_p
pythonapi.PyCapsule_GetPointer.argtypes = [py_object, c_char_p]


class bbox(ctypes.Structure):
    _fields_ = [
        ("xmin", ctypes.c_uint32),
        ("ymin", ctypes.c_uint32),
        ("xmax", ctypes.c_uint32),
        ("ymax", ctypes.c_uint32),
        ("score", ctypes.c_float),
        ("id", ctypes.c_int),
        ("label", ctypes.c_char * 64),
    ]


class Box(ctypes.Structure):
    _fields_ = [("boxesnum", ctypes.c_uint16), ("boxes", ctypes.ARRAY(bbox, 128))]


class ImageInfo(ctypes.Structure):
    _fields_ = [
        ("w", ctypes.c_int),
        ("h", ctypes.c_int),
    ]


class CropImagesInfo(ctypes.Structure):
    _fields_ = [("images", ctypes.ARRAY(ImageInfo, 128))]


class MKey(ctypes.Structure):
    _fields_ = [("num", ctypes.c_int), ("data", ctypes.ARRAY(ctypes.c_char, 5))]


class MKeys(ctypes.Structure):
    _fields_ = [("num", ctypes.c_int), ("data", ctypes.ARRAY(MKey, 7000))]


def check_face_box(face_box):
    w = face_box.xmax - face_box.xmin
    h = face_box.ymax - face_box.ymin

    if w <= 20 or h <= 20:
        return False

    return True


def lyn_free_callback(args):
    for i in args:
        error_check(sdk.lyn_free(i))
    return 0


def return_buffers(callback_data) -> int:
    for buffer in callback_data[1]:
        callback_data[0].push(buffer)
    return 0


def read_keys(keys_file_path, mkeys_dev):
    with open(keys_file_path) as f:
        lines = f.readlines()
    mkeys = MKeys()
    mkeys.num = len(lines)
    for i in range(len(lines)):
        line = lines[i]

        mkeys.data[i].data = line.replace("\n", "").encode("utf-8")
        mkeys.data[i].num = len(mkeys.data[i].data)

    mkeys_np = np.frombuffer(
        mkeys,
        dtype=np.uint8,
        count=ctypes.sizeof(MKeys) // np.dtype(np.uint8).itemsize,
    )

    mkeys_host_ptr = sdk.lyn_numpy_to_ptr(mkeys_np)

    error_check(
        sdk.lyn_memcpy(
            mkeys_dev,
            mkeys_host_ptr,
            ctypes.sizeof(MKeys),
            sdk.lyn_memcpy_dir_t.ClientToServer,
        )
    )


def main():
    parser = argparse.ArgumentParser(description="OCR Detection and Recognition")

    parser.add_argument(
        "input_image",
        metavar="input_image",
        type=str,
        default="../../../data/pp_ocr_sample/ocr_sample.jpeg",
        help="input jpeg file path",
    )
    parser.add_argument(
        "output_image",
        metavar="output_image",
        type=str,
        default="../../../data/pp_ocr_sample/output_ocr.jpeg",
        help="output jpeg file path",
    )
    parser.add_argument(
        "device_id", metavar="device_id", type=int, default=0, help="lynxi device id"
    )
    parser.add_argument(
        "-m1", 
        "--ocr_detect_model",
        metavar="ocr_detect_model",
        type=str,
        default="../../../model/paddleOCRv3_det_b1_h736_w1280_c3_iuint8_ofloat16_0000/Net_0/",
        help="ocr detect model path",
    )
    parser.add_argument(
        "-p", 
        "--ocr_detect_post_plugin",
        metavar="ocr_detect_post_plugin",
        type=str,
        default="../../../common/plugin/postprocess/cross_compiled_lib/libPPOcrDetectPlugin.so",
        help="ocr detect post process plugin path",
    )
    parser.add_argument(
        "-m2", 
        "--ocr_recog_model",
        metavar="ocr_recog_model",
        type=str,
        default="../../../model/paddleOCRv3_rec_b1_h48_w320_c3_iuint8_ofloat16_0000/Net_0",
        help="ocr recog model path",
    )

    parser.add_argument(
        "-m3", 
        "--ocr_cls_model",
        metavar="ocr_cls_model",
        type=str,
        default="../../../model/paddleOCR_cls_b1_h48_w320_c3_iuint8_ofloat16_0000/Net_0",
        help="ocr cls model path",
    )

    parser.add_argument(
        "-o", 
        "--osd_plugin",
        metavar="osd_plugin",
        type=str,
        default="../../../common/plugin/osd/lib/libOsdPlugin.so",
        help="osd plugin path",
    )

    parser.add_argument(
        "-k", 
        "--keys_file_path",
        metavar="keys_file_path",
        type=str,
        default="../../../data/pp_ocr_sample/ppocr_keys_v1.txt",
        help="ppocr keys file path",
    )

    args = parser.parse_args()

    input_file = args.input_image
    output_file = args.output_image
    device_id = args.device_id
    ocr_detect_model = args.ocr_detect_model
    ocr_detect_post_plugin = args.ocr_detect_post_plugin
    ocr_recog_model = args.ocr_recog_model
    ocr_cls_model = args.ocr_cls_model
    osd_plugin_path = args.osd_plugin
    keys_file_path = args.keys_file_path

    if not os.path.exists(input_file):
        raise ("error: input jpeg file does not exist.")

    if not is_valid_device(device_id):
        raise ("error: device id is invalid.")

    if not os.path.exists(ocr_detect_model):
        raise ("error: ocr detect model does not exist.")

    if not os.path.exists(ocr_detect_post_plugin):
        raise ("error: ocr detect post process plugin does not exist.")

    if not os.path.exists(ocr_recog_model):
        raise ("error: ocr feature model does not exist.")

    if not os.path.exists(ocr_cls_model):
        raise ("error: ocr feature model does not exist.")

    if not os.path.exists(osd_plugin_path):
        raise ("error: osd plugin does not exist.")

    if not os.path.exists(keys_file_path):
        raise ("error: keys file does not exist.")
    
    # 2. 创建 context 与 stream
    ctx, ret = sdk.lyn_create_context(device_id)
    error_check(ret, "create context")
    ret = sdk.lyn_register_error_handler(default_stream_error_handle)
    error_check(ret, "lyn_register_error_handler")
    stream, ret = sdk.lyn_create_stream()
    error_check(ret, "lyn_create_stream")

    mkeys_dev, ret = sdk.lyn_malloc(ctypes.sizeof(MKeys))
    common.error_check(ret)
    read_keys(keys_file_path, mkeys_dev)
    mkeys_ptr = pythonapi.PyCapsule_GetPointer(mkeys_dev, None)

    # 3. 初始化解码图片类
    print("image decode")

    image_dec = image_decoder.ImageDecoder()
    image_dec.init(input_file, False)
    imgInfo = image_dec.get_image_dec_info()

    # 4. 同步解码图片到 Device 侧
    decodeImg = sdk.lyn_codec_buf_t()
    decodeImg.size = imgInfo.output.predict_buf_size
    p_decode_buf_out, ret = sdk.lyn_malloc(decodeImg.size)
    error_check(ret)
    decodeImg.data = p_decode_buf_out
    image_dec.decode_image_to_device(stream, decodeImg.data)

    # 5. 加载模型
    ocrDetectModelInfo = ModelInfo(ocr_detect_model)
    ocrRecogModelInfo = ModelInfo(ocr_recog_model)
    ocrClsModelInfo = ModelInfo(ocr_cls_model)

    # 加载lynxi plugin后处理和osd
    ocrPostPlugin, ret = sdk.lyn_plugin_register(ocr_detect_post_plugin)
    common.error_check(ret, "lyn_plugin_register")
    osdPlugin, ret = sdk.lyn_plugin_register(osd_plugin_path)
    common.error_check(ret, "lyn_plugin_register")

    # 6. 检测IPE处理
    print("detect ocr")
    pDetectIpeBuf, ret = sdk.lyn_malloc(ocrDetectModelInfo.input_size)
    error_check(ret)
    ipe_detect = ipe_param_ocr.IpeParamDetect(
        ocrDetectModelInfo.width, ocrDetectModelInfo.height
    )
    ipe_detect.set_img_info(
        imgInfo.output.width,
        imgInfo.output.height,
        sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
    )
    ipe_detect.calc_param(stream, decodeImg.data, pDetectIpeBuf)

    # 7. 调用 APU 推理接口
    pDevDetectApuBuf, ret = sdk.lyn_malloc(ocrDetectModelInfo.output_size)
    error_check(ret)
    error_check(
        sdk.lyn_execute_model_async(
            stream,
            ocrDetectModelInfo.model,
            pDetectIpeBuf,
            pDevDetectApuBuf,
            ocrDetectModelInfo.batch_size,
        )
    )

    # 8.后处理
    pDevDetectApuBuf_ptr = pythonapi.PyCapsule_GetPointer(pDevDetectApuBuf, None)

    pDevDetectInfo, ret = sdk.lyn_malloc(ctypes.sizeof(Box))
    error_check(ret)
    pDevDetectInfo_ptr = pythonapi.PyCapsule_GetPointer(pDevDetectInfo, None)

    crop_images_dev, ret = sdk.lyn_malloc(
        imgInfo.output.width * imgInfo.output.height * 3
    )
    error_check(ret)
    crop_images_dev_ptr = pythonapi.PyCapsule_GetPointer(crop_images_dev, None)

    image_data_dev_ptr = pythonapi.PyCapsule_GetPointer(decodeImg.data, None)

    pDevCropInfo, ret = sdk.lyn_malloc(ctypes.sizeof(CropImagesInfo))
    error_check(ret)
    pDevCropInfo_ptr = pythonapi.PyCapsule_GetPointer(pDevCropInfo, None)

    post_info = struct.pack(
        "4i5P",
        imgInfo.output.height,
        imgInfo.output.width,
        ocrDetectModelInfo.height,
        ocrDetectModelInfo.width,
        pDevDetectApuBuf_ptr,
        pDevDetectInfo_ptr,
        image_data_dev_ptr,
        crop_images_dev_ptr,
        pDevCropInfo_ptr,
    )

    error_check(
        sdk.lyn_plugin_run_async(
            stream,
            ocrPostPlugin,
            "lynPaddleOcrDetectPostProcess",
            post_info,
            len(post_info),
        )
    )

    pHostDetectInfo_np = np.zeros(ctypes.sizeof(Box), dtype=np.uint8)
    pHostDetectInfo = sdk.lyn_numpy_to_ptr(pHostDetectInfo_np)
    error_check(
        sdk.lyn_memcpy_async(
            stream,
            pHostDetectInfo,
            pDevDetectInfo,
            ctypes.sizeof(Box),
            sdk.lyn_memcpy_dir_t.ServerToClient,
        )
    )

    pHostCropInfo_np = np.zeros(ctypes.sizeof(CropImagesInfo), dtype=np.uint8)
    pHostCropInfo = sdk.lyn_numpy_to_ptr(pHostCropInfo_np)
    error_check(
        sdk.lyn_memcpy_async(
            stream,
            pHostCropInfo,
            pDevCropInfo,
            ctypes.sizeof(CropImagesInfo),
            sdk.lyn_memcpy_dir_t.ServerToClient,
        )
    )

    error_check(sdk.lyn_synchronize_stream(stream))

    pHostDetectInfo_c = pythonapi.PyCapsule_GetPointer(pHostDetectInfo, None)
    c_box_info = ctypes.cast(pHostDetectInfo_c, ctypes.POINTER(Box)).contents

    pHostCropInfo_c = pythonapi.PyCapsule_GetPointer(pHostCropInfo, None)
    c_crop_info = ctypes.cast(pHostCropInfo_c, ctypes.POINTER(CropImagesInfo)).contents

    offset = 0

    # cls and recog ipe
    recog_ipe_buffer_pool = bufferpool.buffer_pool(ocrRecogModelInfo.input_size, 5)
    recog_out_buffer_pool = bufferpool.buffer_pool(ocrRecogModelInfo.output_size, 5)
    cls_out_buffer_pool = bufferpool.buffer_pool(ocrClsModelInfo.output_size, 5)

    ipe_recog = ipe_param_ocr.IpeParamRecog(
        ocrRecogModelInfo.width, ocrRecogModelInfo.height
    )

    osd_data = Box()
    osd_data.boxesnum = c_box_info.boxesnum
    for i in range(c_box_info.boxesnum):
        osd_data.boxes[i].xmax = c_box_info.boxes[i].xmax
        osd_data.boxes[i].xmin = c_box_info.boxes[i].xmin
        osd_data.boxes[i].ymax = c_box_info.boxes[i].ymax
        osd_data.boxes[i].ymin = c_box_info.boxes[i].ymin
        osd_data.boxes[i].score = c_box_info.boxes[i].score

        crop_imag_addr = sdk.lyn_addr_seek(crop_images_dev, offset)

        recog_ipe_buffer = recog_ipe_buffer_pool.pop()

        ipe_recog.set_img_info(
            c_crop_info.images[i].w,
            c_crop_info.images[i].h,
            sdk.lyn_pixel_format_t.LYN_PIX_FMT_BGR24,
        )

        ipe_recog.calc_param(stream, crop_imag_addr, recog_ipe_buffer)

        # recog
        recog_out_buffer = recog_out_buffer_pool.pop()

        error_check(
            sdk.lyn_execute_model_async(
                stream,
                ocrRecogModelInfo.model,
                recog_ipe_buffer,
                recog_out_buffer,
                ocrRecogModelInfo.batch_size,
            )
        )

        recog_out_buffer_ptr = pythonapi.PyCapsule_GetPointer(recog_out_buffer, None)
        recog_post_info = struct.pack(
            "3i3P",
            i,
            ocrRecogModelInfo.model_desc.outputTensorAttrArray[0].dims[1],
            ocrRecogModelInfo.model_desc.outputTensorAttrArray[0].dims[2],
            mkeys_ptr,
            recog_out_buffer_ptr,
            pDevDetectInfo_ptr,
        )
        error_check(
            sdk.lyn_plugin_run_async(
                stream,
                ocrPostPlugin,
                "lynPaddleOcrRecogPostProcess",
                recog_post_info,
                len(recog_post_info),
            )
        )

        # cls
        cls_out_buffer = cls_out_buffer_pool.pop()

        error_check(
            sdk.lyn_execute_model_async(
                stream,
                ocrClsModelInfo.model,
                recog_ipe_buffer,
                cls_out_buffer,
                ocrClsModelInfo.batch_size,
            )
        )

        cls_out_host_np = np.zeros(ocrClsModelInfo.output_size // 2, dtype=np.int16)
        cls_out_host = sdk.lyn_numpy_to_ptr(cls_out_host_np)
        error_check(
            sdk.lyn_memcpy_async(
                stream,
                cls_out_host,
                cls_out_buffer,
                ocrClsModelInfo.output_size,
                sdk.lyn_memcpy_dir_t.ServerToClient,
            )
        )

        error_check(sdk.lyn_synchronize_stream(stream))

        pHostRecogInfo_np = np.zeros(ctypes.sizeof(Box), dtype=np.uint8)
        pHostRecogInfo = sdk.lyn_numpy_to_ptr(pHostRecogInfo_np)
        error_check(
            sdk.lyn_memcpy_async(
                stream,
                pHostRecogInfo,
                pDevDetectInfo,
                ctypes.sizeof(Box),
                sdk.lyn_memcpy_dir_t.ServerToClient,
            )
        )

        error_check(sdk.lyn_synchronize_stream(stream))

        pHostRecogInfo_c = pythonapi.PyCapsule_GetPointer(pHostRecogInfo, None)
        c_recog_info = ctypes.cast(pHostRecogInfo_c, ctypes.POINTER(Box)).contents

        if cls_out_host_np[0] < cls_out_host_np[1]:
            label = (
                "text:" + c_recog_info.boxes[i].label.decode("utf-8") + ", angle:180"
            )
        else:
            label = "text:" + c_recog_info.boxes[i].label.decode("utf-8") + ", angle:0"

        osd_data.boxes[i].label = label.encode("utf-8")

        offset += c_crop_info.images[i].w * c_crop_info.images[i].h * 3

        error_check(
            sdk.lyn_stream_add_callback(
                stream,
                return_buffers,
                [recog_ipe_buffer_pool, [recog_ipe_buffer]],
            )
        )

        error_check(
            sdk.lyn_stream_add_callback(
                stream,
                return_buffers,
                [recog_out_buffer_pool, [recog_out_buffer]],
            )
        )

        error_check(
            sdk.lyn_stream_add_callback(
                stream,
                return_buffers,
                [cls_out_buffer_pool, [cls_out_buffer]],
            )
        )

    if osd_data.boxesnum > 0:
        osd_data_np = np.frombuffer(
            osd_data,
            dtype=np.uint8,
            count=ctypes.sizeof(Box) // np.dtype(np.uint8).itemsize,
        )
        osd_data_ptr = sdk.lyn_numpy_to_ptr(osd_data_np)
        osd_data_dev, ret = sdk.lyn_malloc(ctypes.sizeof(Box))
        error_check(
            sdk.lyn_memcpy_async(
                stream,
                osd_data_dev,
                osd_data_ptr,
                ctypes.sizeof(Box),
                sdk.lyn_memcpy_dir_t.ClientToServer,
            )
        )

        # 9. 加载osd Plugin, 将结果叠加到图像上
        print("osd")
        osd_data_dev_ptr = pythonapi.PyCapsule_GetPointer(osd_data_dev, None)
        decodeImg_data_ptr = pythonapi.PyCapsule_GetPointer(decodeImg.data, None)
        osd_para = struct.pack(
            "1P3i1P4i",
            osd_data_dev_ptr,
            imgInfo.output.width,
            imgInfo.output.height,
            sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
            decodeImg_data_ptr,
            24,
            4,
            2,
            4,
        )
        error_check(
            sdk.lyn_plugin_run_async(
                stream, osdPlugin, "lynDrawBoxAndText", osd_para, len(osd_para)
            )
        )

    # 9. 进行图片编码
    error_check(sdk.lyn_synchronize_stream(stream))
    print("save new image")
    encoder = image_encoder.ImageEncoder()
    encoder.encode_image_and_save(stream, output_file, decodeImg, imgInfo)

    #  10. 销毁资源
    ipe_detect.destory()
    error_check(sdk.lyn_synchronize_stream(stream))
    error_check(sdk.lyn_destroy_stream(stream))
    error_check(sdk.lyn_plugin_unregister(ocrPostPlugin))
    error_check(sdk.lyn_plugin_unregister(osdPlugin))

    error_check(sdk.lyn_free(decodeImg.data))
    error_check(sdk.lyn_free(pDetectIpeBuf))
    error_check(sdk.lyn_free(pDevDetectApuBuf))
    error_check(sdk.lyn_free(mkeys_dev))
    error_check(sdk.lyn_free(crop_images_dev))
    error_check(sdk.lyn_free(pDevCropInfo))
    error_check(sdk.lyn_free(osd_data_dev))
    error_check(sdk.lyn_free(pDevDetectInfo))
    ocrDetectModelInfo.unload_model()
    ocrRecogModelInfo.unload_model()
    ocrClsModelInfo.unload_model()

    error_check(sdk.lyn_destroy_context(ctx))
    print("process over")


if __name__ == "__main__":
    main()
